#include <iostream>
#include <memory>
#include <string>
#include <vector>
#include <thread>

// WebSocket++
#define ASIO_STANDALONE
#include <websocketpp/config/asio_no_tls.hpp>
#include <websocketpp/server.hpp>

// sherpa-onnx
#include "sherpa-onnx/c-api/cxx-api.h"

// WebSocket++ type aliases
using server = websocketpp::server<websocketpp::config::asio>;
using connection_hdl = websocketpp::connection_hdl;

// --- 1. 配置 ---
const std::string ASR_MODEL_DIR = "./sherpa-onnx-sense-voice-zh-en-ja-ko-yue-2024-07-17/";
const std::string MATCHA_ACOUSTIC_MODEL = "./matcha-icefall-zh-baker/model-steps-3.onnx";
const std::string MATCHA_VOCODER_MODEL = "./vocos-22khz-univ.onnx";
const std::string MATCHA_LEXICON_PATH = "./matcha-icefall-zh-baker/lexicon.txt";
const std::string MATCHA_TOKENS_PATH = "./matcha-icefall-zh-baker/tokens.txt";
const std::string MATCHA_DICT_DIR = "./matcha-icefall-zh-baker/dict";
const std::string MATCHA_RULE_FSTS = "./matcha-icefall-zh-baker/phone.fst,./matcha-icefall-zh-baker/date.fst,./matcha-icefall-zh-baker/number.fst";

// 全局对象
std::unique_ptr<sherpa_onnx::cxx::OfflineTts> g_tts;
std::shared_ptr<sherpa_onnx::cxx::OfflineRecognizer> global_recognizer;

// --- 2. WebSocket 事件处理 ---
void on_open(connection_hdl hdl)
{
    std::cout << "Isla客户端已连接!" << std::endl;
}

void on_close(connection_hdl hdl)
{
    std::cout << "Isla客户端已断开连接。" << std::endl;
}

void on_message(server *s, connection_hdl hdl, server::message_ptr msg)
{
    if (msg->get_opcode() == websocketpp::frame::opcode::text)
    {
        // 文本转语音
        std::string text = msg->get_payload();
        auto audio = g_tts->Generate(text, 0, 1.2);
        s->send(hdl, audio.samples.data(), audio.samples.size() * sizeof(float), websocketpp::frame::opcode::binary);
        std::cout << "已发送TTS音频。" << std::endl;
    }
    else
    {
        // 语音识别转文本
        const auto *samples = reinterpret_cast<const float *>(msg->get_payload().data());
        int32_t num_samples = msg->get_payload().size() / sizeof(float);

        // 创建识别流
        sherpa_onnx::cxx::OfflineStream stream = global_recognizer->CreateStream();
        stream.AcceptWaveform(16000, samples, num_samples);
        global_recognizer->Decode(&stream);

        std::string result = global_recognizer->GetResult(&stream).text;
        s->send(hdl, result, websocketpp::frame::opcode::text);
        std::cout << "已发送识别结果: " << result << std::endl;
    }
}

// --- 3. 主函数 ---
int main()
{
    // --- ASR配置（sense-voice模型） ---
    sherpa_onnx::cxx::OfflineRecognizerConfig offline_config;
    offline_config.model_config.sense_voice.model = ASR_MODEL_DIR + "model.int8.onnx";
    offline_config.model_config.sense_voice.language = "auto";
    offline_config.model_config.sense_voice.use_itn = false;
    offline_config.model_config.tokens = ASR_MODEL_DIR + "tokens.txt";
    offline_config.model_config.num_threads = 2;
    offline_config.model_config.debug = false;

    global_recognizer = std::make_shared<sherpa_onnx::cxx::OfflineRecognizer>(
        sherpa_onnx::cxx::OfflineRecognizer::Create(offline_config));

    // --- TTS配置（MatchaTTS）---
    sherpa_onnx::cxx::OfflineTtsConfig tts_config;
    tts_config.model.matcha.acoustic_model = MATCHA_ACOUSTIC_MODEL;
    tts_config.model.matcha.vocoder = MATCHA_VOCODER_MODEL;
    tts_config.model.matcha.lexicon = MATCHA_LEXICON_PATH;
    tts_config.model.matcha.tokens = MATCHA_TOKENS_PATH;
    tts_config.model.matcha.dict_dir = MATCHA_DICT_DIR;
    tts_config.model.num_threads = 1;
    tts_config.model.debug = false;
    tts_config.rule_fsts = MATCHA_RULE_FSTS;

    try
    {
        g_tts = std::make_unique<sherpa_onnx::cxx::OfflineTts>(
            sherpa_onnx::cxx::OfflineTts::Create(tts_config));
    }
    catch (const std::exception &e)
    {
        std::cerr << "初始化TTS失败: " << e.what() << std::endl;
        return -1;
    }

    // --- 启动WebSocket服务器 ---
    server ws_server;
    try
    {
        ws_server.init_asio();
        ws_server.set_open_handler(&on_open);
        ws_server.set_close_handler(&on_close);
        ws_server.set_message_handler(bind(&on_message, &ws_server, std::placeholders::_1, std::placeholders::_2));

        ws_server.listen(8765);
        ws_server.start_accept();

        std::cout << "启动实时语音对话服务器..." << std::endl;
        ws_server.run();
    }
    catch (const std::exception &e)
    {
        std::cerr << "服务器错误: " << e.what() << std::endl;
    }

    return 0;
}